from neo4j import GraphDatabase
import time
import os


class Neo4j_DB:

    def __init__(self, uri="neo4j://152.136.185.123:9099", user="neo4j", password="digitalmofs", name='module-kg.db'):
        self.driver = GraphDatabase.driver(uri, auth=(user, password), name=name)

    def close(self):
        self.driver.close()

    def send_message_to_db(self, message):
        job, model = self.prase_info(message)
        if job['operation'] == 'load':
            find = self.find_relation(model['model_id'], job['id'], 'LOAD')
            if not find:
                self.load_model(job, model)
        elif job['operation'] == 'save':
            find = self.find_relation(model['model_id'], job['id'], 'SAVE')
            if not find:
                self.save_model(job, model)

    def prase_info(self, message):
        job = {
            'md5_time_ms': message['md5_time_ms'],
            'io_time_ms': message['io_time_ms'],
            'path': message['path'],
            'jobid': message['jobid'],
            'username': message['username'],
            'cluster': message['cluster'],
            'operation': message['operation'],
            'timestamp_millis': message['timestamp_millis'],
        }
        if job['jobid'] is None:
            job['id'] = job['cluster'] + '-' + str(os.getpid())
        else:
            job['id'] = job['cluster'] + '-' + job['jobid']

        model = {
            'model_id': message['model']['model_hash'],
            'model_hash': message['model']['model_hash'],
            'model_short_hash': message['model']['model_short_hash'],
            'model_size_bytes': message['model']['model_size_bytes'],
        }

        return job, model

    def load_model(self, job, model=None,):
        with self.driver.session() as session:
            session.write_transaction(self._load_model, job, model)

    def save_model(self, job, model):
        with self.driver.session() as session:
            session.write_transaction(self._save_model, job, model)

    @staticmethod
    def _load_model(tx, job, model):
        query = (
            "MERGE (job:Job { id: $id,\
                    username: $username, \
                    cluster: $cluster}) "
            "MERGE (model:Model { model_id:$model_id, model_hash: $model_hash, \
                    model_short_hash: $model_short_hash, model_size_bytes: $model_size_bytes}) "
            "CREATE (model)-[:LOAD {path: $path, io_time_ms: $io_time_ms, \
                md5_time_ms: $md5_time_ms, time_stamp: $time_stamp}]->(job) "
            "RETURN job, model"
        )
        time_stamp = job['timestamp_millis'] / 1000.0
        time_stamp = time.strftime(
            "%Y-%m-%d %H:%M:%S", time.localtime(time_stamp))
        tx.run(query,
               id=job['id'], md5_time_ms=job['md5_time_ms'], io_time_ms=job['io_time_ms'],
               path=job['path'], username=job['username'],
               cluster=job['cluster'], model_hash=model['model_hash'],
               model_id=model['model_id'], model_size_bytes=model['model_size_bytes'],
               model_short_hash=model['model_short_hash'], time_stamp=time_stamp)

    @staticmethod
    def _save_model(tx, job, model):
        query = (
            "MERGE (job:Job { id: $id, \
                    username: $username, \
                    cluster: $cluster}) "
            "MERGE (model:Model { model_id:$model_id, model_hash: $model_hash, \
                    model_short_hash: $model_short_hash, model_size_bytes: $model_size_bytes}) "
            "CREATE (job)-[:SAVE {path: $path, io_time_ms: $io_time_ms, \
                md5_time_ms: $md5_time_ms, time_stamp: $time_stamp}]->(model) "
            "RETURN job, model"
        )
        time_stamp = job['timestamp_millis'] / 1000.0
        time_stamp = time.strftime(
            "%Y-%m-%d %H:%M:%S", time.localtime(time_stamp))
        tx.run(query,
               id=job['id'], md5_time_ms=job['md5_time_ms'], io_time_ms=job['io_time_ms'],
               path=job['path'], username=job['username'],
               cluster=job['cluster'], model_hash=model['model_hash'],
               model_id=model['model_id'], model_size_bytes=model['model_size_bytes'],
               model_short_hash=model['model_short_hash'], time_stamp=time_stamp)

    def find_relation(self, id, model_id, opration):
        with self.driver.session() as session:
            find_relation = session.write_transaction(
                self._find_relation, id, model_id, opration)
            return find_relation

    @staticmethod
    def _find_relation(tx, model_id, id, opration):
        if opration == "LOAD":
            query = (
                "MATCH (model:Model{ model_id:$model_id})-[relation:LOAD]->(job:Job{ id:$id})"
                "RETURN relation"
            )
        else:
            query = (
                "MATCH (job:Job{ id:$id})-[relation:SAVE]->(model:Model{ model_id:$model_id})"
                "RETURN relation"
            )
        result = tx.run(query, model_id=model_id, id=id)
        find_relation = result.single()
        return find_relation

    def update_model_kestrel(self, model_meta, job_info):
        with self.driver.session() as session:
            session.write_transaction(self._update_model_kestrel, model_meta, job_info)

    @staticmethod
    def _update_model_kestrel(tx, model_meta, job_info):
        model_hash, model_short_hash, model_size_bytes = model_meta
        username = job_info['username']
        cluster = job_info['cluster']
        if job_info['jobid'] is None:
            job_id = job_info['cluster'] + '-' + str(os.getpid())
        else:
            job_id = job_info['cluster'] + '-' + job_info['jobid']
        time_stamp = job_info['timestamp_millis'] / 1000.0
        time_stamp = time.strftime(
            "%Y-%m-%d %H:%M:%S", time.localtime(time_stamp))
        path = os.getcwd()
        query = (
            "MERGE (job:Job { id: $id,\
                    username: $username, \
                    cluster: $cluster}) "
            "MERGE (model:Model { model_id:$model_id, model_hash: $model_hash, \
                    model_short_hash: $model_short_hash, model_size_bytes: $model_size_bytes}) "
            "SET model.to_kestrel = true "
            "CREATE (job)-[:TOKESTREL { path: $path, time_stamp: $time_stamp}]->(model) "
            "RETURN model"
        )
        tx.run(query, model_id=model_hash, model_hash=model_hash,
               model_short_hash=model_short_hash, model_size_bytes=model_size_bytes,
               id=job_id, username=username, cluster=cluster, time_stamp=time_stamp, path=path)
